#include "detector.h"

SparkDetector::SparkDetector() {}

cv::Mat SparkDetector::preprocessColor(const cv::Mat& image) {
    // 颜色检测
    cv::Mat hsv, mask1, mask2, colorMask;

    // 转换为HSV颜色空间
    cv::cvtColor(image, hsv, cv::COLOR_BGR2HSV);

    // 定义火光的HSV范围（红色）
    cv::Scalar lower1(0, 120, 230);
    cv::Scalar upper1(60, 255, 255);

    // cv::Scalar lower2(170, 120, 70);
    // cv::Scalar upper2(180, 255, 255);

    // 创建掩码
    cv::inRange(hsv, lower1, upper1, mask1);
    //cv::inRange(hsv, lower2, upper2, mask2);

    // 合并掩码
    //cv::bitwise_or(mask1, mask2, colorMask);
    colorMask = mask1.clone();

    // 形态学操作
    cv::Mat kernel = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(5, 5));
    cv::morphologyEx(colorMask, colorMask, cv::MORPH_CLOSE, kernel);
    cv::morphologyEx(colorMask, colorMask, cv::MORPH_OPEN, kernel);

    cv::dilate(colorMask, colorMask, kernel, cv::Point(-1, -1), 1);

    return colorMask;
}

cv::Mat SparkDetector::preprocessMotion(const cv::Mat& image) {
    if (previousFrame.empty()) {
        return cv::Mat::zeros(image.size(), CV_8UC1);
    }

    cv::Mat grayFrame, grayPrevious;
    cvtColor(image, grayFrame, cv::COLOR_BGR2GRAY);
    cvtColor(previousFrame, grayPrevious, cv::COLOR_BGR2GRAY);

    cv::Mat diff;
    absdiff(grayPrevious, grayFrame, diff);

    // 阈值处理
    cv::Mat motionMask;
    threshold(diff, motionMask, 30, 255, cv::THRESH_BINARY);

    // 形态学操作
    cv::Mat kernel = getStructuringElement(cv::MORPH_ELLIPSE, cv::Point(5, 5));
    morphologyEx(motionMask, motionMask, cv::MORPH_CLOSE, kernel);

    return motionMask;
}

cv::Mat SparkDetector::preprocessImage(const cv::Mat& image) {
    // 颜色检测
    cv::Mat colorMask, motionMask, combinedMask;

    colorMask = preprocessColor(image);
    if (!previousFrame.empty()) {
        motionMask = preprocessMotion(image);
    }

    if (!motionMask.empty()) {
        combinedMask = colorMask & motionMask;
    } else {
        combinedMask = colorMask;
    }

    previousFrame = image.clone();

    return combinedMask;
}

std::vector<cv::Rect> SparkDetector::postprocessRegions(std::vector<std::vector<cv::Point>>& contours) {
    std::vector<cv::Rect> sparkRegions;

    for (auto& contour : contours) {
        double area = cv::contourArea(contour);
        if (area > 20) { // 忽略小区域
            cv::Rect rect = cv::boundingRect(contour);
            sparkRegions.push_back(rect);
        }
    }

    return sparkRegions;
}

std::vector<cv::Rect> SparkDetector::detect(const cv::Mat& image) {
    cv::Mat mask = preprocessImage(image);

    // 查找轮廓
    std::vector<std::vector<cv::Point>> contours;
    cv::findContours(mask, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

    std::vector<cv::Rect> rects = postprocessRegions(contours);
    return rects;
}

cv::Mat SparkDetector::annotateImage(const cv::Mat& image, const std::vector<cv::Rect>& sparkRegions) {
    cv::Mat annotated = image.clone();

    for (const auto& rect : sparkRegions) {
        // 绘制矩形框
        cv::rectangle(annotated, rect, cv::Scalar(0, 0, 255), 2);
    }

    return annotated;
}

PassenerDetector::PassenerDetector() {}

std::vector<cv::Rect> PassenerDetector::detect(const cv::Mat& image) {
    std::vector<cv::Rect> faces;

    cv::CascadeClassifier faceCascade;
    faceCascade.load(".\\haarcascade_frontalface_default.xml");
    if (faceCascade.empty()) { std::cout << "XML file not loaded" << std::endl; }

    faceCascade.detectMultiScale(image, faces, 1.1, 10);
    return faces;
}

cv::Mat PassenerDetector::annotateImage(const cv::Mat& image, const std::vector<cv::Rect>& Regions) {
    cv::Mat annotated = image.clone();
    for (const auto& rect : Regions) {
        // 绘制矩形框
        cv::rectangle(annotated, rect, cv::Scalar(0, 255, 0), 2);
    }
    return annotated;
}

IntruderDetector::IntruderDetector(double lr, int thresh, int minArea)
    : learningRate(lr), thresholdValue(thresh), minContourArea(minArea) {}

std::vector<cv::Rect> IntruderDetector::detect(const cv::Mat& image) {
    std::vector<cv::Rect> intruders;

    // 转换为灰度图
    cv::Mat gray;
    cvtColor(image, gray, cv::COLOR_BGR2GRAY);

    // 高斯模糊，减少噪声影响
    GaussianBlur(gray, gray, cv::Point(21, 21), 0);

    // cv::imshow("gray",gray);
    // cv::waitKey(10);

    // 如果背景为空，初始化背景
    if (background.empty()) {
        gray.convertTo(background,CV_32F);
        return intruders;
    }

    cv::Mat grayFloat;
    gray.convertTo(grayFloat, CV_32F);

    // 计算当前帧与背景的差异
    cv::Mat diff, diffFloat;
    absdiff(background, grayFloat, diffFloat);

    diffFloat.convertTo(diff, CV_8U);

    cv::imshow("diff",diff);
    cv::waitKey(10);

    // 二值化差异图像
    cv::Mat thresh;
    threshold(diff, thresh, thresholdValue, 255, cv::THRESH_BINARY);

    // 形态学操作，去除噪声并填充空洞
    cv::Mat kernel = getStructuringElement(cv::MORPH_ELLIPSE, cv::Point(5, 5));
    morphologyEx(thresh, thresh, cv::MORPH_CLOSE, kernel);
    morphologyEx(thresh, thresh, cv::MORPH_OPEN, kernel, cv::Point(-1,-1), 2);

    cv::imshow("thresh",thresh);
    cv::waitKey(10);
    // 查找轮廓
    std::vector<std::vector<cv::Point>> contours;
    findContours(thresh, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
    // 处理每个轮廓
    for (const auto& contour : contours) {
        // 忽略小面积的轮廓
        if (contourArea(contour) < minContourArea) {
            continue;
        }

        // 获取轮廓的边界矩形
        cv::Rect rect = boundingRect(contour);
        intruders.push_back(rect);
    }

    // 更新背景模型
    accumulateWeighted(grayFloat, background, learningRate);
    return intruders;
}

cv::Mat IntruderDetector::annotateImage(const cv::Mat& image, const std::vector<cv::Rect>& Regions) {
    cv::Mat annotated = image.clone();
    //std::cout<<Regions.size()<<std::endl;
    for (const auto& rect : Regions) {
        // 绘制矩形框
        cv::rectangle(annotated, rect, cv::Scalar(255, 0, 0), 2);
    }
    return annotated;
}
